"""
实验结果收集器
用于收集和上传实验数据
"""
import os
import json
import zipfile
import shutil
import requests
import io
from typing import Dict, Any, Optional, List, Union
from datetime import datetime
from pathlib import Path
from urllib.parse import urljoin
import logging

try:
    import pandas as pd
    HAS_PANDAS = True
except ImportError:
    HAS_PANDAS = False

try:
    import pyarrow as pa
    import pyarrow.parquet as pq
    HAS_PYARROW = True
except ImportError:
    HAS_PYARROW = False

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .client import ParameterClient

from .exceptions import ParameterClientError

logger = logging.getLogger(__name__)


class ExperimentCollector:
    """实验结果收集器"""
    
    def __init__(self, client: "ParameterClient", name: Optional[str] = None, description: Optional[str] = None):
        """
        创建实验收集器
        
        Args:
            client: ParameterClient实例
            name: 实验名称（可选，默认自动生成）
            description: 实验描述（可选）
        """
        if not client.project:
            raise ValueError("项目未加载，请先调用 load_project()")
        
        self.client = client
        self.project = client.project
        self.project_name_en = client.project_name
        
        # 生成实验名称
        self.name = name or f"实验_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.description = description
        
        # 本地存储根目录（使用时间戳作为目录名）
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        self.local_root = Path(".local_results") / self.project_name_en / timestamp
        self.restore_dir = self.local_root / "restore"
        self.view_dir = self.local_root / "view"
        
        # 创建目录结构
        self._setup_directories()
        
        # 保存参数快照（如果还没有保存）
        self._save_snapshot_if_needed()
        
        # 数据存储
        self.categories: Dict[str, Dict[str, Any]] = {}  # category_name -> {description, sort_order}
        self.scalar_data: List[Dict[str, Any]] = []  # 单值数据列表
        self.series_data: List[Dict[str, Any]] = []  # 序列数据元信息列表
        self.result_files: List[Dict[str, Any]] = []  # 结果文件列表
        
        logger.info(f"创建实验收集器: {self.name}")
    
    def _setup_directories(self):
        """创建本地存储目录结构"""
        self.restore_dir.mkdir(parents=True, exist_ok=True)
        self.view_dir.mkdir(parents=True, exist_ok=True)
        (self.restore_dir / "files").mkdir(exist_ok=True)
        (self.restore_dir / "series").mkdir(exist_ok=True)
        (self.view_dir / "files").mkdir(exist_ok=True)
        (self.view_dir / "series").mkdir(exist_ok=True)
        (self.view_dir / "scalar").mkdir(exist_ok=True)
    
    def _save_snapshot_if_needed(self):
        """保存参数快照到实验目录"""        
        # 如果有快照数据（从服务器加载），直接dump成JSON
        if self.client.snapshot_data:
            snapshot_path = self.local_root / "snapshot.json"
            restore_snapshot_path = self.restore_dir / "snapshot.json"
            
            with open(snapshot_path, 'w', encoding='utf-8') as f:
                json.dump(self.client.snapshot_data, f, ensure_ascii=False, indent=2)
            
            with open(restore_snapshot_path, 'w', encoding='utf-8') as f:
                json.dump(self.client.snapshot_data, f, ensure_ascii=False, indent=2)
            
            logger.info(f"已保存参数快照: {snapshot_path}")
            return
    
    def create_category(self, name: str, description: Optional[str] = None):
        """
        创建结果分类（幂等操作）
        
        Args:
            name: 分类名称
            description: 分类描述（可选）
        """
        if name in self.categories:
            # 已存在，不重复创建
            logger.debug(f"分类已存在: {name}")
            return
        
        sort_order = len(self.categories)
        self.categories[name] = {
            "name": name,
            "description": description,
            "sort_order": sort_order
        }
        logger.info(f"创建分类: {name}")
    
    def add_scalar(self, category: str, name: str, value: Any, 
                   unit: Optional[str] = None, description: Optional[str] = None, 
                   data_group: Optional[str] = None):
        """
        添加单值数据
        
        Args:
            category: 分类名称
            name: 指标名称
            value: 指标值（可以是数字、字符串等）
            unit: 单位（可选）
            description: 描述（可选）
            data_group: 数据分组（可选）
        """
        # 确保分类存在
        self.create_category(category)
        
        # 检查是否已存在（同一分类下指标名唯一）
        for item in self.scalar_data:
            if item["category"] == category and item["name"] == name:
                raise ValueError(f"指标 '{name}' 在分类 '{category}' 下已存在")
        
        self.scalar_data.append({
            "category": category,
            "name": name,
            "value": value,
            "unit": unit,
            "description": description,
            "data_group": data_group or None,
            "created_at": datetime.now().isoformat()
        })
        
        logger.info(f"添加单值数据: {category}/{name} = {value}")
    
    def add_series(self, category: str, name: str, 
                   values: List[Any] = None,
                   index: Optional[List[Any]] = None,
                   series: Optional[Any] = None,  # 支持pandas Series
                   unit: Optional[str] = None,
                   description: Optional[str] = None,
                   data_group: Optional[str] = None):
        """
        添加序列数据
        
        Args:
            category: 分类名称
            name: 指标名称
            values: 值列表（必需，除非提供了series）
            index: 索引列表（可选），如果不提供则自动生成 [1, 2, 3, ...]
            unit: 单位（可选）
            description: 描述（可选）
            data_group: 数据分组（可选）
        """
        if values is None or len(values) == 0:
            raise ValueError("values不能为空")
        
        # 确保分类存在
        self.create_category(category)
        
        # 检查是否已存在（同一分类下指标名唯一）
        for item in self.series_data:
            if item["category"] == category and item["name"] == name:
                raise ValueError(f"指标 '{name}' 在分类 '{category}' 下已存在")
        
        # 处理索引
        if index is None:
            # 自动生成索引 [1, 2, 3, ...]
            index = list(range(1, len(values) + 1))
        else:
            if len(index) != len(values):
                raise ValueError(f"索引长度({len(index)})与值长度({len(values)})不匹配")
        
        # 保存为Parquet文件
        if not HAS_PYARROW:
            raise ValueError("需要安装pyarrow以保存序列数据")
        
        # 创建分类目录
        category_dir = self.restore_dir / "series" / category
        category_dir.mkdir(parents=True, exist_ok=True)
        category_view_dir = self.view_dir / "series" / category
        category_view_dir.mkdir(parents=True, exist_ok=True)
        
        # 保存Parquet文件
        parquet_path = category_dir / f"{name}.parquet"
        df_data = {
            "index": index,
            "value": values
        }
        df = pd.DataFrame(df_data)
        df.to_parquet(parquet_path, index=False)
        
        # 保存CSV文件（用于查看）
        csv_path = category_view_dir / f"{name}.csv"
        df.to_csv(csv_path, index=False, encoding='utf-8-sig')
        
        # 保存元信息
        self.series_data.append({
            "category": category,
            "name": name,
            "parquet_path": f"series/{category}/{name}.parquet",
            "unit": unit,
            "description": description,
            "data_group": data_group or None,
            "index_column_name": "index",
            "data_length": len(values),
            "created_at": datetime.now().isoformat()
        })
        
        logger.info(f"添加序列数据: {category}/{name} (长度: {len(values)})")
    
    def upload_file(self, name: str, file_path: str, 
                    file_type: str = "other",
                    description: Optional[str] = None):
        """
        上传结果文件
        
        Args:
            name: 文件名称
            file_path: 本地文件路径
            file_type: 文件类型（image, excel, csv, other）
            description: 文件描述（可选）
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        file_size = os.path.getsize(file_path)
        
        # 复制文件到restore/files和view/files
        restore_file_path = self.restore_dir / "files" / os.path.basename(file_path)
        view_file_path = self.view_dir / "files" / os.path.basename(file_path)
        
        shutil.copy2(file_path, restore_file_path)
        shutil.copy2(file_path, view_file_path)
        
        # 保存元信息
        self.result_files.append({
            "name": name,
            "file_path": f"files/{os.path.basename(file_path)}",
            "file_type": file_type,
            "file_size": file_size,
            "description": description,
            "created_at": datetime.now().isoformat()
        })
        
        logger.info(f"上传文件: {name} ({file_type}, {file_size} bytes)")
    
    def _generate_metadata(self) -> Dict[str, Any]:
        """生成metadata.json"""
        metadata = {
            "experiment": {
                "name": self.name,
                "description": self.description,
                "project_name_en": self.project_name_en,
                "created_at": datetime.now().isoformat()
            },
            "results": self.result_files,
            "categories": [
                {
                    "name": cat["name"],
                    "description": cat.get("description"),
                    "sort_order": cat.get("sort_order", 0)
                }
                for cat in sorted(self.categories.values(), key=lambda x: x.get("sort_order", 0))
            ],
            "data_series": [
                {
                    "name": item["name"],
                    "category_name": item["category"],
                    "data_group": item.get("data_group"),
                    "parquet_path": item["parquet_path"],
                    "unit": item.get("unit"),
                    "description": item.get("description"),
                    "index_column_name": item["index_column_name"],
                    "data_length": item["data_length"],
                    "created_at": item.get("created_at")
                }
                for item in self.series_data
            ],
            "data_scalar": [
                {
                    "name": item["name"],
                    "category_name": item["category"],
                    "data_group": item.get("data_group"),
                    "value": item["value"],
                    "unit": item.get("unit"),
                    "description": item.get("description"),
                    "created_at": item.get("created_at")
                }
                for item in self.scalar_data
            ]
        }
        return metadata
    
    def _save_scalar_view_files(self):
        """保存单值数据到view/scalar/目录（按分类组织）"""
        # 按分类组织
        by_category: Dict[str, List[Dict[str, Any]]] = {}
        for item in self.scalar_data:
            category = item["category"]
            if category not in by_category:
                by_category[category] = []
            by_category[category].append({
                "name": item["name"],
                "value": item["value"],
                "unit": item.get("unit"),
                "description": item.get("description")
            })
        
        # 保存JSON文件
        for category, items in by_category.items():
            json_path = self.view_dir / "scalar" / f"{category}.json"
            with open(json_path, 'w', encoding='utf-8') as f:
                json.dump(items, f, ensure_ascii=False, indent=2)
    
    def _check_server_connectivity(self) -> bool:
        """
        检查服务器是否可连通
        
        Returns:
            bool: 如果服务器可连通返回True，否则返回False
        """
        try:
            # 使用健康检查接口来测试连通性（健康检查接口在根路径，不在/api下）
            health_url = f"http://{self.client.host}:{self.client.port}/health"
            # 设置较短的超时时间，用于快速检测连通性
            response = requests.get(health_url, timeout=5)
            response.raise_for_status()
            logger.info(f"服务器连通性检查成功: {self.client.host}:{self.client.port}")
            return True
        except requests.exceptions.ConnectionError:
            logger.warning(f"无法连接到服务器 {self.client.host}:{self.client.port}，将仅保存本地文件")
            return False
        except requests.exceptions.Timeout:
            logger.warning(f"连接服务器 {self.client.host}:{self.client.port} 超时，将仅保存本地文件")
            return False
        except Exception as e:
            logger.warning(f"检查服务器连通性时出错: {str(e)}，将仅保存本地文件")
            return False
    
    def commit(self):
        """提交所有收集的数据到服务器（如果服务器可连通）"""
        if not self.project_name_en:
            raise ValueError("项目名称未设置")
        
        # 如果描述为空，提示用户输入
        if not self.description:
            description = input("请输入实验描述（直接回车跳过）: ").strip()
            if description:
                self.description = description
        
        logger.info(f"开始提交实验: {self.name}")
        
        # 保存单值数据到view目录
        self._save_scalar_view_files()
        
        # 生成metadata.json（使用更新后的description）
        metadata = self._generate_metadata()
        metadata_path = self.restore_dir.parent / "metadata.json"
        with open(metadata_path, 'w', encoding='utf-8') as f:
            json.dump(metadata, f, ensure_ascii=False, indent=2)
        
        # 创建ZIP文件（包含restore目录和metadata.json）
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        zip_path = self.restore_dir.parent / f"upload_{timestamp}.zip"
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
            # 添加restore目录下的所有文件（保持目录结构）
            for root, dirs, files in os.walk(self.restore_dir):
                for file in files:
                    file_path = Path(root) / file
                    # 计算相对于restore_dir的路径，作为ZIP内的路径
                    arc_name = file_path.relative_to(self.restore_dir)
                    # 在ZIP中保存为 restore/... 路径
                    zf.write(file_path, f"restore/{arc_name}")
            
            # 添加metadata.json
            zf.write(metadata_path, "metadata.json")
        
        logger.info(f"实验数据已保存到本地: {self.local_root}")
        
        # 检查服务器连通性，只有在连通时才上传
        if not self._check_server_connectivity():
            logger.info("服务器不可用，实验数据已保存到本地，未上传到服务器")
            logger.info(f"本地数据路径: {self.local_root}")
            # 不清理ZIP文件，保留作为备份
            return {
                "experiment_id": None,
                "uploaded": False,
                "local_path": str(self.local_root),
                "message": "服务器不可用，数据已保存到本地"
            }
        
        # 服务器可连通，尝试上传
        endpoint = f"/projects/{self.project_name_en}/experiments/upload"
        
        # 准备multipart/form-data请求（只上传ZIP文件，metadata在ZIP中）
        # 确保endpoint不以/开头，使用与_make_request相同的URL构建方式
        if endpoint.startswith('/'):
            endpoint = endpoint[1:]
        url = urljoin(self.client.base_url + '/', endpoint)
        
        # 直接使用文件对象上传
        upload_success = False
        try:
            logger.info(f"准备上传文件到: {url}")
            logger.info(f"ZIP文件路径: {zip_path}, 文件大小: {zip_path.stat().st_size} bytes")
            
            # 打开文件并上传
            with open(zip_path, 'rb') as file_obj:
                files = {
                    'file': (f"restore_{timestamp}.zip", file_obj, 'application/zip')
                }
                
                # 准备headers：只保留认证相关的headers，移除Content-Type
                headers = {}
                for key, value in self.client.session.headers.items():
                    if key.lower() not in ['content-type']:
                        headers[key] = value
                
                logger.debug(f"请求headers: {headers}")
                logger.debug(f"files参数: file字段已准备")
                
                # 直接使用requests.post而不是session.post，避免session的默认设置干扰
                response = requests.post(
                    url,
                    files=files,
                    headers=headers,
                    timeout=300  # 设置较长的超时时间，因为文件可能较大
                )
                response.raise_for_status()
            
            result = response.json()
            logger.info(f"实验提交成功: {result.get('experiment_id', 'N/A')}")
            result["uploaded"] = True
            result["local_path"] = str(self.local_root)
            upload_success = True
            return result
        except requests.exceptions.HTTPError as e:
            # 记录详细的错误信息
            error_detail = ""
            try:
                error_detail = e.response.text
                logger.error(f"HTTP错误 {e.response.status_code}: {error_detail}")
            except:
                pass
            logger.warning(f"上传失败，但数据已保存到本地: {self.local_root}")
            # 上传失败时不抛出异常，只记录警告，数据已经保存在本地
            return {
                "experiment_id": None,
                "uploaded": False,
                "local_path": str(self.local_root),
                "error": f"上传失败: {str(e)}",
                "message": "上传失败，但数据已保存到本地"
            }
        except Exception as e:
            logger.error(f"提交实验失败: {str(e)}")
            logger.warning(f"上传失败，但数据已保存到本地: {self.local_root}")
            # 上传失败时不抛出异常，只记录警告，数据已经保存在本地
            return {
                "experiment_id": None,
                "uploaded": False,
                "local_path": str(self.local_root),
                "error": str(e),
                "message": "上传失败，但数据已保存到本地"
            }
        finally:
            # 只有上传成功时才清理临时ZIP文件，否则保留作为备份
            if upload_success and zip_path.exists():
                try:
                    zip_path.unlink()
                    logger.debug(f"已清理临时ZIP文件: {zip_path}")
                except Exception:
                    pass
    
    def upload(self):
        """commit的别名"""
        return self.commit()

